# encoding=utf-8
from ultralytics import YOLO
import torch
# 加载预训练模型
model = YOLO("yolov8n.pt")    # 加载预先训练的模型（建议用于训练）
# Use the model
if __name__ == '__main__':
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    # results = model.train(data='datasets/TrafficSignData/data.yaml', epochs=350, batch=4)  # 训练模型
    results = model.train(data='datasets/TrafficSignData/data.yaml', epochs=500, batch=5)  # 训练模型
    # 将模型转为onnx格式,能够看到模型的详细结构
    # success = model.export(format='onnx')



